import logging
from datetime import datetime
from pinecone import Pinecone, ServerlessSpec
from sklearn.datasets import load_digits
from tqdm import tqdm
import numpy as np
from collections import Counter

# 设置日志格式，包括日期和时间
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S')

pinecone = Pinecone(api_key="1569e82f-e6cb-45ef-9de2-2cdc5c74809d")
index_name = "mnist-index"

# 创建索引
existing_indexes = pinecone.list_indexes()
if any(index['name'] == index_name for index in existing_indexes):
    logging.info(f"索引 '{index_name}' 已存在，正在删除...")
    pinecone.delete_index(index_name)
    logging.info(f"索引 '{index_name}' 已成功删除。")
else:
    logging.info(f"索引 '{index_name}' 不存在，将创建新索引。")

logging.info(f"正在创建新索引 '{index_name}'...")
pinecone.create_index(
    name=index_name,
    dimension=64,
    metric="euclidean",
    spec=ServerlessSpec(
        cloud="aws",
        region="us-east-1"
    )
)
logging.info(f"索引 '{index_name}' 创建成功。")

# 连接到索引
index = pinecone.Index(index_name)
logging.info(f"已成功连接到索引 '{index_name}'。")

# 加载MNIST数据集
digits = load_digits(n_class=10)
X = digits.data
y = digits.target

# 划分训练集和测试集
train_size = int(0.8 * len(X))
X_train, X_test = X[:train_size], X[train_size:]
y_train, y_test = y[:train_size], y[train_size:]

# 准备数据并上传到Pinecone
vectors = []
for i in range(len(X_train)):
    vector_id = str(i)
    vector_values = X_train[i].tolist()
    metadata = {"label": int(y_train[i])}
    vectors.append((vector_id, vector_values, metadata))

batch_size = 1000
for i in range(0, len(vectors), batch_size):
    batch = vectors[i:i + batch_size]
    index.upsert(batch)
    logging.info(f"已上传 {i + len(batch)} 条数据到索引。")

# 测试k=11时的准确率
correct_count = 0
total_count = 0
for i in tqdm(range(len(X_test))):
    test_vector = X_test[i].tolist()
    results = index.query(vector=test_vector, top_k=11, include_metadata=True)
    labels = [match['metadata']['label'] for match in results['matches']]
    predicted_label = Counter(labels).most_common(1)[0][0]
    if predicted_label == y_test[i]:
        correct_count += 1
    total_count += 1

accuracy = correct_count / total_count
logging.info(f"当k=11时，使用Pinecone的准确率为：{accuracy:.4f}")
